from fastapi import FastAPI, HTTPException, Depends
from pydantic import BaseModel
from celery.result import AsyncResult
from sqlalchemy.orm import Session
from app import tasks  # Fixed import
from app import models  # Fixed import
from app.database import engine, get_db  # Fixed import
import json

# Create all tables
models.Base.metadata.create_all(bind=engine)

app = FastAPI(title="FastAPI with Celery and Redis")

class TaskInput(BaseModel):
    number: int


@app.get("/")
async def root():
    return {"message": "Hello World! This is FastAPI with Celery and Redis"}


@app.post("/tasks/add")
async def add_task(task_input: TaskInput, db: Session = Depends(get_db)):
    """
    Create a task to add numbers from 0 to the input number
    """
    # Start the Celery task
    task = tasks.add_numbers.delay(task_input.number)
    
    # Create a new task record in the database
    db_task = models.Task(
        id=task.id,
        status="PENDING",
        task_name="add_numbers",
        input_data={"number": task_input.number},
        result={"progress": 0}  # Initialize progress at 0
    )
    db.add(db_task)
    db.commit()
    
    return {"task_id": task.id, "message": f"Task created to sum numbers from 0 to {task_input.number}"}


@app.get("/tasks/{task_id}")
async def get_task_status(task_id: str, db: Session = Depends(get_db)):
    """
    Get the status of a task
    """
    # Get the task from the database
    db_task = db.query(models.Task).filter(models.Task.id == task_id).first()
    
    if not db_task:
        raise HTTPException(status_code=404, detail="Task not found")
    
    # Get the latest status from Celery
    task_result = AsyncResult(task_id)
    
    # Update the task status in the database if it has changed
    celery_status = task_result.status
    if db_task.status != celery_status:
        if celery_status == "PENDING":
            db_task.status = "PENDING"
        elif celery_status == "STARTED":
            db_task.status = "PROCESSING"
        elif celery_status == "SUCCESS":
            db_task.status = "DONE"
            if task_result.ready():
                db_task.result = task_result.get()
        else:
            db_task.status = celery_status
            
        db.commit()
    
    # Prepare the response
    result = {
        "task_id": task_id,
        "status": db_task.status,
        "created_at": db_task.created_at,
        "updated_at": db_task.updated_at,
        "task_name": db_task.task_name,
        "input_data": db_task.input_data,
        "progress": db_task.result.get("progress", 0) if db_task.result else 0
    }
    
    # Add the full result if the task has been completed
    if db_task.result and "result" in db_task.result:
        result["result"] = db_task.result
    
    return result


@app.get("/tasks")
async def list_tasks(db: Session = Depends(get_db)):
    """
    Get a list of all tasks
    """
    db_tasks = db.query(models.Task).all()
    
    return [
        {
            "task_id": task.id,
            "status": task.status,
            "task_name": task.task_name,
            "created_at": task.created_at,
            "updated_at": task.updated_at,
            "progress": task.result.get("progress", 0) if task.result else 0
        }
        for task in db_tasks
    ]


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)